import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
import numpy as np
import pandas as pd
from langchain_experimental.tools import PythonAstREPLTool

if __name__ == '__main__':
    load_dotenv(override=True)

    df = pd.read_csv("archive/WA_Fn-UseC_-Telco-Customer-Churn.csv")
    pd.set_option('max_colwidth', 200)
    # print(df.head(5))
    tool = PythonAstREPLTool(locals={"df": df})
    result = tool.invoke("df['SeniorCitizen'].mean()")
    print(result)
    print(df['SeniorCitizen'].mean())
    print(df['MonthlyCharges'].mean())
    print(tool.invoke("df['MonthlyCharges'].mean()"))
    # prompt = ChatPromptTemplate.from_messages(
    #     [("system", "你叫做小智，是一个乐于助人的助手，请根据用户的输入的问题进行回答"),
    #      ("human", "{input}")],
    # )
    #
    # DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    # model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    #
    # parser = StrOutputParser()
    # chain = prompt | model | parser
